import paddle
import numpy as np
import matplotlib.pyplot as plt
import meshio
import pickle
import pandas as pd
from scipy.spatial import Delaunay
from meshpy.tet import MeshInfo, build, Options

# def standardize_data(data):
#     mean = paddle.mean(data, axis=0)  
#     stddev = paddle.std(data, axis=0)  
#     return (data - mean) / stddev  , mean, stddev

# 转化从CFD得到的坐标和标签为paddle张量，用于训练，返回坐标和对应的CFD标签（速度、压力、密度）
def to_tensor(xyz_and_labels):
    data_points = xyz_and_labels
    data_xyz = data_points[:, 0:3]
    
    mean = np.mean(data_xyz, axis=0)  
    std = np.std(data_xyz, axis=0)  
    
    # 标准化  
    data_xyz = (data_xyz - mean) / std  
    
    data_p = data_points[:, 3:4]
    data_u = data_points[:, 4:5]
    data_v = data_points[:, 5:6]
    data_w = data_points[:, 6:7]

    x_pde_train = paddle.to_tensor(data_xyz, dtype='float32')
    # x_pde_train = x_pde_train
    x_pde_train.stop_gradient = False #计算并传播梯度。默认为True。
    
    p_label = paddle.to_tensor(data_p, dtype='float32')
    u_label = paddle.to_tensor(data_u, dtype='float32')
    v_label = paddle.to_tensor(data_v, dtype='float32')
    w_label = paddle.to_tensor(data_w, dtype='float32')


    label = {}
    label['u'] = u_label
    label['v'] = v_label
    label['w'] = w_label
    label['p'] = p_label

    return x_pde_train, label


# 对残差进行排序，从大到小，并选取前k个
def find_ids_of_top_k(array, k):
    ids =  sorted(range(len(array)), key=lambda i: array[i], reverse=True)[:k] 
    return np.array(ids)


# 展示所有的顶点
def show_points(points, size = 1,color="C0", marker="*", dpi = 100, figsize = (20,10)):
    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    ax.set_aspect('equal')
    ax.scatter(points[:, 0], points[:, 1], s=size, color=color, marker=marker)
    plt.show()


def find_center_in_tetrahedron(tetrahedrons, vertices):
    centroids = []
    for tetrahedron in tetrahedrons:
        vertex1 = vertices[tetrahedron[0]]
        vertex2 = vertices[tetrahedron[1]]
        vertex3 = vertices[tetrahedron[2]]
        vertex4 = vertices[tetrahedron[3]]
        centroid = (vertex1 + vertex2 + vertex3 + vertex4) / 4
        centroids.append(centroid)
    centroids_points = np.array(centroids)    
    return centroids_points


def tetrahedron_midpoint(tetrahedrons, vertices):
    all_points_in_half_edge = []
    for tetrahedron in tetrahedrons:
        nodes_01 = vertices[tetrahedron[0]]
        nodes_02 = vertices[tetrahedron[1]]
        nodes_03 = vertices[tetrahedron[2]]
        nodes_04 = vertices[tetrahedron[3]]
        midpoint12 = (nodes_01 + nodes_02) / 2
        all_points_in_half_edge.append(midpoint12)
        midpoint13 = (nodes_01 + nodes_03) / 2
        all_points_in_half_edge.append(midpoint13)
        midpoint14 = (nodes_01 + nodes_04) / 2
        all_points_in_half_edge.append(midpoint14)
        midpoint23 = (nodes_02 + nodes_03) / 2
        all_points_in_half_edge.append(midpoint23)
        midpoint24 = (nodes_02 + nodes_04) / 2
        all_points_in_half_edge.append(midpoint24)
        midpoint34 = (nodes_03 + nodes_04) / 2
        all_points_in_half_edge.append(midpoint34)
    all_points_in_half_edge = np.array(all_points_in_half_edge)    
    return all_points_in_half_edge

# def remove_duplicate_points(points):
#     unique_points = np.unique(points, axis=0)
#     return unique_points

def plot_room_add_points_3d(points):
    fig = plt.figure(figsize=(10,10))
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(points[:, 0], points[:, 2], points[:, 1], c='red',s=5,marker='*')
    # ax.scatter(points[:, 0], points[:, 1], points[:, 2], c='k',s=5,marker='*')
    ax.set_xlabel('X')
    ax.set_ylabel('Z')
    ax.set_zlabel('Y')  
    ax.set_xlim(0,3)
    ax.set_ylim(0,3)
    ax.set_zlim(0,3)
    # ax.invert_zaxis()  # 调整y轴朝上
    # ax.view_init(elev=30, azim=0)  # 设置视角，使y轴朝上
    # plt.axis("scaled")
    line_points_1 = np.array([[3, 0, 3], [3, 3, 3]])  # 定义线段的起点和终点坐标  
    ax.plot(line_points_1[:, 0], line_points_1[:, 1], line_points_1[:, 2], c='k')  
    line_points_2 = np.array([[0, 0, 3], [3, 0, 3]])  # 定义线段的起点和终点坐标  
    ax.plot(line_points_2[:, 0], line_points_2[:, 1], line_points_2[:, 2], c='k')  
    line_points_3 = np.array([[3, 0, 3], [3, 0, 0]])  # 定义线段的起点和终点坐标  
    ax.plot(line_points_3[:, 0], line_points_3[:, 1], line_points_3[:, 2], c='k') 
    plt.show()

# def process_mesh(mesh):
#     points_from_mesh = mesh.points
#     # points = utils.remove_duplicate_points(points_from_mesh)
#     processed_points = np.unique(points_from_mesh, axis=0)
#     tri = Delaunay(points, furthest_site=False, incremental=False, qhull_options=None)

#     mesh_info = MeshInfo()
#     mesh_info.set_points(tri.points)
#     mesh_info.set_facets(tri.convex_hull)
#     mesh_info.set_elements(tri.simplices)
#     mesh = build(mesh_info, options=Options("p"))

#     nodes = []
#     for i, n in enumerate(mesh.points):
#         nodes.append(n)
#     nodes = np.array(nodes, dtype=int)

#     cells = []
#     for j, c in enumerate(mesh.elements):
#         cells.append(c)
#     cells = np.array(cells)

#     # save vtk file
#     vtkFileName = "work/progress/level_01/mesh_files/level_01.vtk"
#     mesh.write_vtk(vtkFileName)

def process_mesh(mesh):
    points_from_mesh = mesh.points
    processed_points = np.unique(points_from_mesh, axis=0)
    tri = Delaunay(processed_points, furthest_site=False, incremental=False, qhull_options=None)

    inf = MeshInfo()
    inf.set_points(tri.points)
    inf.set_facets(tri.convex_hull)
    inf.set_elements(tri.simplices)
    mesh = build(inf, options=Options("p"))
    nodes = []
    for i, n in enumerate(mesh.points):
        nodes.append(n)
    nodes = np.array(nodes, dtype=int)
    cells = []
    for j, c in enumerate(mesh.elements):
        cells.append(c)
    cells = np.array(cells)
    # save vtk file
    vtkFileName = "work/progress/level_01/mesh_files/level_01.vtk"
    mesh.write_vtk(vtkFileName)


'''

def plot_loss(level='level_01',Adam_epoches=50000):
    data = np.loadtxt(f'work/progress/{level}/loss/loss1.txt')
    data2 = np.loadtxt(f'work/progress/{level}/loss/loss2.txt')
    # epoch = data[:, 0]
    # loss_train = data[:, 1]
    fig, ax = plt.subplots(figsize=(7,4),dpi=150)
    plt.plot(data[:, 0], data[:, 1], label=r'Adam')
    plt.plot(Adam_epoches+data2[10:, 0], data2[10:, 1], label=r'LBFGS')
    plt.legend()
    plt.yscale('log')
    plt.ylabel('Loss')
    plt.xlabel('Iterations')

def plot_drag_lift():
    
    pinns_ret = pd.read_csv('work/evaluation/lift_drag_compare/pinns.csv')
    fluent_ret = pd.read_csv('work/evaluation/lift_drag_compare/cfd.csv')
    # nodes = pinns_ret['Nodes'] 
    cells = pinns_ret['Cells'] 
    drag = pinns_ret['Drag'] 
    lift = pinns_ret['Lift']
    fig,ax = plt.subplots(2,1,figsize=(7,4),dpi=300,sharex=True)
    ax[0].plot(cells, drag,'-P',c='k',label='Pinns')
    ax[1].plot(cells, lift,'-P',c='k',label='Pinns')
    ax[0].plot(fluent_ret['Cells'], fluent_ret['Drag'],'-o',c='skyblue',label='CFD')
    ax[1].plot(fluent_ret['Cells'], fluent_ret['Lift'],'-o',c='skyblue',label='CFD')
    ax[0].set_ylabel(r'Drag')
    ax[1].set_ylabel(r'Lift')
    ax[1].set_xlabel(r'Cells')
    # ax[0].set_ylim([2000,3000])
    # ax[1].set_ylim([22000,32000])
    ax[0].legend(loc=4)
    ax[1].legend(loc=4)
# 原始顶点和三角网格质心的可视化
def show_center_and_vertices(central_points, delaunay_mesh):
    fig, ax = plt.subplots(figsize=(20,8))
    ax.set_aspect('equal')
    plt.scatter(central_points[:, 0], central_points[:, 1],s=1,c='red',label='Center of mass points')
    plt.scatter(delaunay_mesh['vertices'][:,0], delaunay_mesh['vertices'][:,1],s=1,label='Vertexes')
    plt.legend()
    plt.show()


def process_index(array):
    array_list = [list(arr) for arr in array]
    flattened_list = [item for sublist in array_list for item in sublist]
    unique_list = list(set(flattened_list))
    return unique_list

def get_internal_points(arr1, *arrs):
    mask = np.ones(len(arr1), dtype=bool)
    for i, row in enumerate(arr1):
        for compare_row in np.concatenate(arrs):
            if np.array_equal(row, compare_row):
                mask[i] = False
                break
    result = arr1[mask]
    return result

# 计算网格的面积，适合三角形
def calculate_triangle_area(vertices, triangles):
    """
    计算每个三角形的面积

    参数:
        vertices: 三角形网格的顶点坐标 (N, 2)
        triangles: 三角形的索引 (M, 3)

    返回:
        包含每个三角形面积的NumPy数组，形状为 (M, 1)
    """
    areas = np.zeros((triangles.shape[0], 1))

    for i, triangle in enumerate(triangles):
        # 获取三角形的顶点坐标
        p1 = vertices[triangle[0]]
        p2 = vertices[triangle[1]]
        p3 = vertices[triangle[2]]

        # 计算三角形的面积
        area = 0.5 * np.abs(np.cross(p2 - p1, p3 - p1))

        # 将面积存入数组
        areas[i] = area

    return areas



# 对网格进行前处理，需要锚定网格的翼型边界、远场矩形边界
def process_mesh(path):
    #处理翼型点，并排序
    def uniform_points(coordinates):
        center_x = coordinates[:, 0].mean()
        center_y = coordinates[:, 1].mean()
        angle = np.arctan2(coordinates[:, 1] - center_y, coordinates[:, 0] - center_x)
        points_indices = np.argsort(angle)
        ordered_points = coordinates[points_indices, :]
        return ordered_points
    # 从meshio读入文件，以ansys形式
    mesh = meshio.read(path, file_format='ansys')
    # 得到cells,并获得index
    cells = mesh.cells
    edges = cells[0].data
    farfield_edges = np.concatenate((cells[1].data, cells[2].data, cells[3].data, cells[4].data), axis=0)
    airfoil_edges = cells[5].data
    # 得到节点的index
    all_points_index = process_index(edges)
    farfield_points_index = process_index(farfield_edges)
    airfoil_points_index = process_index(airfoil_edges)
    # 得到顶点坐标 
    vertices = mesh.points[all_points_index]
    farfield_points = mesh.points[farfield_points_index]
    airfoil_points_0 = mesh.points[airfoil_points_index]
    internal_points = get_internal_points(vertices, farfield_points, airfoil_points_0)

    #
    airfoil_points = np.vstack((airfoil_points_0[1:104,:], airfoil_points_0[0:1,:], airfoil_points_0[104:,:]))
    farfield_points = uniform_points(farfield_points)

    # 合并所有得到的顶点，用于生成网格
    polygon = np.concatenate((airfoil_points, farfield_points, internal_points))
    # 设置约束边界的索引，翼型边界和矩形边界
    inner_airfoil_index = np.hstack((np.arange(0, airfoil_points.shape[0], 1).reshape(airfoil_points.shape[0], 1), 
                                    np.arange(1, airfoil_points.shape[0]+1, 1).reshape(airfoil_points.shape[0], 1)))
    inner_airfoil_index[airfoil_points.shape[0]-1, 1]=0
    outer_farfield_index = np.hstack((np.arange(0, farfield_points.shape[0], 1).reshape(farfield_points.shape[0], 1), 
                                    np.arange(1, farfield_points.shape[0] +1, 1).reshape(farfield_points.shape[0], 1)))
    outer_farfield_index[farfield_points.shape[0]-1, 1]=0
    # outer_farfield_indexd的每个索引从 inner_airfoil_index.shape[0]=205开始计数，即[205,...]
    outer_farfield_index = outer_farfield_index + inner_airfoil_index.shape[0] 
     #所有边界顺序排列，形成闭环，注意，这里一定要形成闭环
    inner_outer_index = np.concatenate((inner_airfoil_index, outer_farfield_index), axis=0)
    # return all_points, farfield_points, airfoil_points, internal_points
    return polygon, inner_outer_index

# 对残差进行排序，从大到小，并选取前k个
def find_ids_of_top_k(array, k):
    ids =  sorted(range(len(array)), key=lambda i: array[i], reverse=True)[:k] 
    return np.array(ids)

# 读取字典
def load_dict(file_path):
    with open(file_path, 'rb') as file:
        dictionary = pickle.load(file)
    return dictionary
# 保存字典
def save_dict(dictionary, file_path):
    with open(file_path, 'wb') as file:
        pickle.dump(dictionary, file)
# 保存网格
def save_mesh(C, path):
    # 保存字典到文件
    save_dict(C, path + '.pickle')
    #输出网格的vtk以及msh文件
    # 获取三角剖分结果的顶点和单元格信息
    vertices = C['vertices']
    cells = [("triangle", C['triangles'])]
    # 创建meshio的网格对象
    mesh = meshio.Mesh(points=vertices, cells=cells)
    # 定义保存的文件名
    vtkFileName = path + '.vtk'
    # 将网格对象保存为VTK文件
    meshio.write(vtkFileName, mesh, file_format="vtk", binary=False)
    # 使用meshio读取VTK文件
    mesh = meshio.read(vtkFileName)
    # 将网格对象保存为Gmsh支持的格式（例如.msh文件）
    mshFileName = path + '.msh'
    meshio.write(mshFileName, mesh, file_format='gmsh22',binary=False)
    
'''